#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
医患对话数据分析
本地运行：python generate.py
接口：
  GET  /           -> 前端页面
  GET  /api/dialogs -> 返回分析结果
  GET  /api/dialogs.csv -> 返回CSV下载
"""
import pandas as pd
import numpy as np
import jieba
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report
from transformers import BertTokenizer, BertForSequenceClassification
import torch

# --------------------------
# 1. 数据加载与预处理（适配多科室对话）
# --------------------------
def load_and_preprocess_data(data_path):
    """
    加载医患对话数据，预处理格式：[科室, 医生对话, 患者对话, 对话标签(0:良好,1:紧张)]
    """
    # 示例数据格式（实际替换为你的数据路径，支持CSV/Excel）
    df = pd.read_csv(data_path)
    # 1.1 合并医患对话（用于情感分析和关系预测）
    df["full_dialog"] = df["doctor_dialog"] + " " + df["patient_dialog"]
    # 1.2 中文分词（适配BERT和TF-IDF）
    df["cut_dialog"] = df["full_dialog"].apply(lambda x: " ".join(jieba.lcut(x.strip())))
    # 1.3 按科室分组（后续针对性推荐沟通技巧）
    dept_groups = df.groupby("department").agg({"full_dialog": list, "label": list}).to_dict("index")
    return df, dept_groups

# --------------------------
# 2. 情感分析模块（基于BERT中文情感模型）
# --------------------------
class DialogSentimentAnalyzer:
    def __init__(self):
        # 加载预训练中文情感分析模型（适配医疗场景，可替换为自定义训练模型）
        self.tokenizer = BertTokenizer.from_pretrained("hfl/chinese-bert-wwm-ext")
        self.model = BertForSequenceClassification.from_pretrained(
            "hfl/chinese-bert-wwm-ext", num_labels=3  # 3类情感：0-负面，1-中性，2-正面
        )
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)
        self.sentiment_map = {0: "负面（可能存在冲突风险）", 1: "中性（沟通平稳）", 2: "正面（医患信任良好）"}

    def predict_sentiment(self, dialog_text):
        """预测单条对话的情感倾向"""
        inputs = self.tokenizer(
            dialog_text, truncation=True, padding="max_length", max_length=128, return_tensors="pt"
        ).to(self.device)
        self.model.eval()
        with torch.no_grad():
            outputs = self.model(**inputs)
            logits = outputs.logits
            pred_label = torch.argmax(logits, dim=1).item()
        return self.sentiment_map[pred_label]

# --------------------------
# 3. 医患关系预测模块（基于TF-IDF+逻辑回归，轻量化易部署）
# --------------------------
class DoctorPatientRelationPredictor:
    def __init__(self):
        self.vectorizer = TfidfVectorizer(max_features=2000, stop_words=["的", "了", "是", "我", "你"])  # 医疗领域停用词可扩展
        self.classifier = LogisticRegression(class_weight="balanced")  # 解决数据不平衡问题

    def train_model(self, train_texts, train_labels):
        """用标注数据训练关系预测模型（0:良好，1:紧张）"""
        # 文本转TF-IDF特征
        X_train = self.vectorizer.fit_transform(train_texts)
        # 训练分类器
        self.classifier.fit(X_train, train_labels)
        print("模型训练完成！")

    def predict_relation(self, dialog_text):
        """预测单条对话的医患关系状态"""
        X = self.vectorizer.transform([dialog_text])
        pred_label = self.classifier.predict(X)[0]
        pred_prob = self.classifier.predict_proba(X)[0][pred_label]  # 预测概率
        relation_map = {0: f"良好（置信度：{pred_prob:.2f}）", 1: f"紧张（置信度：{pred_prob:.2f}）"}
        return relation_map[pred_label]

    def evaluate_model(self, test_texts, test_labels):
        """评估模型性能（可选，用于验证数据效果）"""
        X_test = self.vectorizer.transform(test_texts)
        y_pred = self.classifier.predict(X_test)
        print("\n=== 医患关系预测模型评估报告 ===")
        print(classification_report(test_labels, y_pred, target_names=["关系良好", "关系紧张"]))

# --------------------------
# 4. 沟通技巧推荐模块（基于科室+情感/关系标签匹配）
# --------------------------
def get_communication_tips(department, sentiment_result, relation_result):
    """
    根据科室、情感倾向、关系状态，推荐针对性沟通技巧
    可扩展：基于实际科室数据补充更多细分场景技巧
    """
    # 基础技巧库（按科室和风险等级分类）
    tips_lib = {
        "内科": {
            "负面/紧张": [
                "1. 先倾听患者主诉（如“您先慢慢说，我都记下来了”），避免打断；",
                "2. 用通俗语言解释病情（如“这个指标高说明炎症，不是严重问题”），减少信息差；",
                "3. 主动告知下一步诊疗计划（如“明天做检查，结果出来我第一时间跟您说”）。"
            ],
            "中性/良好": [
                "1. 保持微笑和眼神交流，强化信任（如“您恢复得不错，继续按方案来”）；",
                "2. 适当提问关注患者感受（如“最近睡眠怎么样？有没有哪里不舒服？”）。"
            ]
        },
        "外科": {
            "负面/紧张": [
                "1. 术前详细说明手术风险（分点列出，避免模糊表述）；",
                "2. 术后主动反馈恢复情况（如“今天伤口换药很顺利，再观察2天就能出院”）；",
                "3. 对患者担忧的疼痛问题，明确告知缓解措施（如“止痛药会按时给您，不会让您疼的”）。"
            ],
            "中性/良好": [
                "1. 用手势辅助解释手术部位（如指向X光片标注位置）；",
                "2. 鼓励患者提问（如“您还有什么不清楚的，随时问我”）。"
            ]
        },
        # 可补充：儿科、妇产科、急诊科等其他科室技巧
        "其他科室": {
            "负面/紧张": [
                "1. 降低沟通语速，避免专业术语堆砌；",
                "2. 重复关键信息确保患者理解（如“我再跟您确认下，药每天吃2次，对吗？”）；",
                "3. 表达共情（如“我理解您现在很着急，我们一起想办法”）。"
            ],
            "中性/良好": [
                "1. 保持沟通节奏，及时回应患者疑问；",
                "2. 结尾给予积极预期（如“按这个方案，您的情况会逐步好转的”）。"
            ]
        }
    }

    # 匹配科室
    target_dept = department if department in tips_lib else "其他科室"
    # 匹配风险等级（负面/紧张 归为高风险，中性/良好 归为低风险）
    risk_level = "负面/紧张" if ("负面" in sentiment_result or "紧张" in relation_result) else "中性/良好"
    # 返回技巧
    return tips_lib[target_dept][risk_level]

# --------------------------
# 5. 主函数：整合所有功能，运行分析
# --------------------------
def main(data_path):
    # 步骤1：加载预处理数据
    df, dept_groups = load_and_preprocess_data(data_path)
    print(f"数据加载完成！共{len(df)}条对话，涉及科室：{list(dept_groups.keys())}")

    # 步骤2：初始化模型
    sentiment_analyzer = DialogSentimentAnalyzer()
    relation_predictor = DoctorPatientRelationPredictor()

    # 步骤3：训练医患关系预测模型（用80%数据训练，20%测试）
    train_texts, test_texts, train_labels, test_labels = train_test_split(
        df["cut_dialog"], df["label"], test_size=0.2, random_state=42
    )
    relation_predictor.train_model(train_texts, train_labels)
    relation_predictor.evaluate_model(test_texts, test_labels)  # 可选：评估模型

    # 步骤4：示例：分析单条对话（可循环批量处理所有数据）
    sample_idx = 0  # 取第一条数据作为示例，可替换为任意索引
    sample_data = df.iloc[sample_idx]
    print(f"\n=== 医患对话分析示例（{sample_data['department']}）===")
    print(f"医生对话：{sample_data['doctor_dialog']}")
    print(f"患者对话：{sample_data['patient_dialog']}")

    # 4.1 情感分析
    sentiment_res = sentiment_analyzer.predict_sentiment(sample_data["full_dialog"])
    print(f"1. 对话情感倾向：{sentiment_res}")

    # 4.2 医患关系预测
    relation_res = relation_predictor.predict_relation(sample_data["cut_dialog"])
    print(f"2. 医患关系预测：{relation_res}")

    # 4.3 沟通技巧推荐
    tips = get_communication_tips(sample_data["department"], sentiment_res, relation_res)
    print(f"3. 针对性沟通技巧：")
    for tip in tips:
        print(f"   {tip}")

    # （可选）批量分析：循环df中所有数据，将结果保存到CSV
    # df["情感倾向"] = df["full_dialog"].apply(sentiment_analyzer.predict_sentiment)
    # df["医患关系预测"] = df["cut_dialog"].apply(relation_predictor.predict_relation)
    # df["沟通技巧"] = df.apply(lambda x: get_communication_tips(x["department"], x["情感倾向"], x["医患关系预测"]), axis=1)
    # df.to_csv("医患对话分析结果.csv", index=False, encoding="utf-8-sig")
    # print("\n批量分析完成！结果已保存至：医患对话分析结果.csv")

# --------------------------
# 运行入口（替换为你的数据路径）
# --------------------------
if __name__ == "__main__":
    # 数据格式要求：CSV文件需包含列名：department, doctor_dialog, patient_dialog, label
    # label说明：0=医患关系良好，1=医患关系紧张（需提前标注部分数据用于训练）
    YOUR_DATA_PATH = "医患对话数据.csv"  # 替换为你的数据文件路径
    main(YOUR_DATA_PATH)